# This file contains modules common to various models


from utils.utils import *


def DWConv(c1, c2, k=1, s=1, act=True):
    # Depthwise convolution
    return Conv(c1, c2, k, s, g=math.gcd(c1, c2), act=act)


class Conv(nn.Module):
    # Standard convolution
    def __init__(self, c1, c2, k=1, s=1, g=1, act=True):  # ch_in, ch_out, kernel, stride, groups
        super(Conv, self).__init__()
        self.conv = nn.Conv2d(c1, c2, k, s, k // 2, groups=g, bias=False)
        self.bn = nn.BatchNorm2d(c2)
        self.act = nn.LeakyReLU(0.1, inplace=True) if act else nn.Identity()

    def forward(self, x):

        x = self.conv(x)
        x = self.bn(x)
        x = self.act(x)
        return x
        # return self.act(self.bn(self.conv(x)))

    def fuseforward(self, x):
        return self.act(self.conv(x))


class Bottleneck(nn.Module):
    # Standard bottleneck
    def __init__(self, c1, c2, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, shortcut, groups, expansion
        super(Bottleneck, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_, c2, 3, 1, g=g)
        self.add = shortcut and c1 == c2

    def forward(self, x):
        return x + self.cv2(self.cv1(x)) if self.add else self.cv2(self.cv1(x))


class BottleneckCSP(nn.Module):
    # CSP Bottleneck https://github.com/WongKinYiu/CrossStagePartialNetworks
    def __init__(self, c1, c2, n=1, shortcut=True, g=1, e=0.5):  # ch_in, ch_out, number, shortcut, groups, expansion
        super(BottleneckCSP, self).__init__()
        c_ = int(c2 * e)  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = nn.Conv2d(c1, c_, 1, 1, bias=False)
        self.cv3 = nn.Conv2d(c_, c_, 1, 1, bias=False)
        self.cv4 = Conv(c2, c2, 1, 1)
        self.bn = nn.BatchNorm2d(2 * c_)  # applied to cat(cv2, cv3)
        self.act = nn.LeakyReLU(0.1, inplace=True)
        self.m = nn.Sequential(*[Bottleneck(c_, c_, shortcut, g, e=1.0) for _ in range(n)])

    def forward(self, x):
        y1 = self.cv3(self.m(self.cv1(x)))
        y2 = self.cv2(x)
        return self.cv4(self.act(self.bn(torch.cat((y1, y2), dim=1))))


class SPP(nn.Module):
    # Spatial pyramid pooling layer used in YOLOv3-SPP
    def __init__(self, c1, c2, k=(5, 9, 13)):
        super(SPP, self).__init__()
        c_ = c1 // 2  # hidden channels
        self.cv1 = Conv(c1, c_, 1, 1)
        self.cv2 = Conv(c_ * (len(k) + 1), c2, 1, 1)
        self.m = nn.ModuleList([nn.MaxPool2d(kernel_size=x, stride=1, padding=x // 2) for x in k])

    def forward(self, x):
        x = self.cv1(x)
        return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1))


class Flatten(nn.Module):
    # Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
    def forward(self, x):
        return x.view(x.size(0), -1)


class Focus(nn.Module):
    # Focus wh information into c-space
    def __init__(self, c1, c2, k=1):
        super(Focus, self).__init__()
        self.conv = Conv(c1 * 4, c2, k, 1)

    def forward(self, x):  # x(b,c,w,h) -> y(b,4c,w/2,h/2)
        return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1))


class Concat(nn.Module):
    # Concatenate a list of tensors along dimension
    def __init__(self, dimension=1):
        super(Concat, self).__init__()
        self.d = dimension

    def forward(self, x):
        return torch.cat(x, self.d)


class yolov5_backbone(nn.Module):

    def __init__(self, input_img_ch = 3,depth_multiple = 0.34, width_multiple = 0.50):
        
        super(yolov5_backbone, self).__init__()

        # self.training_mode = training_mode

        # print ("training_mode:",self.training_mode)

        start_ch = int(64 * width_multiple)

        # 1-P1/2
        self.focus = Focus(input_img_ch, start_ch,3) 
        # 2-P2/4
        self.conv1 = Conv(start_ch, start_ch*2, 3, 2)

        self.bottleneck_csp1 = nn.Sequential(*[Bottleneck(start_ch*2,start_ch*2) for i in range(int(round(3 * depth_multiple)))]) # 3 *0.33
        # 4-P3/8
        self.conv2 = Conv(start_ch*2, start_ch*4, 3, 2)

        # self.bottleneck_csp2 = nn.Sequential(*[Bottleneck(128,128) for i in range(3)]) # 9 *0.33
        self.bottleneck_csp2 = BottleneckCSP(start_ch*4,start_ch*4,int(round(9 * depth_multiple)))

        # 6-P4/16
        self.conv3 = Conv(start_ch*4, start_ch*8, 3, 2)

        # self.bottleneck_csp3 = nn.Sequential(*[Bottleneck(256,256) for i in range(3)]) # 9 *0.33
        self.bottleneck_csp3 = BottleneckCSP(start_ch*8,start_ch*8,int(round(9 * depth_multiple)))

        # 8-P5/32
        self.conv4 = Conv(start_ch*8, start_ch*16, 3, 2)

        self.spp = SPP(start_ch*16,start_ch*16, [5,9,13])

        # 10
        # self.bottleneck_csp4 = nn.Sequential(*[BottleneckCSP(512,512) for i in range(2)]) # 6 *0.33

        self.bottleneck_csp4 = BottleneckCSP(start_ch*16,start_ch*16,int(round(6 * depth_multiple)))
        


    def forward(self, x):# [1,3,640,640]

        # anchors = self.anchors(x)
        
        x_list = []

        x = self.focus(x) # torch.Size([1, 32, 320, 320])
        

        x = self.conv1(x) # torch.Size([1, 64, 160, 160])
        x = self.bottleneck_csp1(x) 
        # x_list.append(x)


        x = self.conv2(x) # torch.Size([1, 128, 80, 80])
        x = self.bottleneck_csp2(x)
        x_list.append(x)
        

        x = self.conv3(x) # torch.Size([1, 256, 40, 40])
        x = self.bottleneck_csp3(x)
        x_list.append(x)
        

        x = self.conv4(x) # torch.Size([1, 512, 20, 20])
        
        x = self.spp(x)
        
        x = self.bottleneck_csp4(x)
        x_list.append(x)

        return x_list

class yolov5_head_without_detect(nn.Module):

    def __init__(self, na_nc_5, nc= 5, strides =[8,16,32],  in_channels=(256, 512, 1024),depth_multiple = 0.34, width_multiple = 0.50):
        super(yolov5_head_without_detect, self).__init__()

        c1 = int(in_channels[-1] * width_multiple)
        n1 = int(3 * depth_multiple)

        self.c3_1 = BottleneckCSP(c1,c1,n1)
        self.conv2d_1 = nn.Conv2d(c1,na_nc_5,1,1,0)

        c2 = int(in_channels[-2] * width_multiple)


        self.upsample_1 = nn.Upsample(scale_factor=2, mode='nearest')
        self.concat_1 = Concat()  
        self.conv_2 = Conv(c2 + c1,c2,1,1)
        self.c3_2 = BottleneckCSP(c2,c2,n1)
        self.conv2d_2 = nn.Conv2d(c2,na_nc_5,1,1,0)


        c3 = int(in_channels[-3] * width_multiple)

        self.upsample_2 = nn.Upsample(scale_factor=2, mode='nearest')
        self.concat_2 = Concat()
        self.conv_3 = Conv(c3 + c2,c3,1,1)
        self.c3_3 = BottleneckCSP(c3,c3,n1)
        self.conv2d_3 = nn.Conv2d(c3,na_nc_5,1,1,0)

        # self. f = [-1,16,11]
        # self. i = 22
        self.na = 3
        self.stride= torch.tensor(strides)
        self.nc = nc

        self._initialize_biases()

    def _initialize_biases(self,cf=None):  

        model_list = [self.conv2d_3, self.conv2d_2, self.conv2d_1]

        for mi,  s in zip(model_list,  self.stride):  #  from

            b = mi.bias.view(self.na, -1)  # conv.bias(255) to (3,85)
            with torch.no_grad():
                b[:, 4] += math.log(8 / (640 / s) ** 2)  # obj (8 objects per 640 image)
                b[:, 5:] += math.log(0.6 / (self.nc - 0.99)) if cf is None else torch.log(cf / cf.sum())  # cls
            mi.bias = torch.nn.Parameter(b.view(-1), requires_grad=True)
      


    def forward(self, x_in):


        x = self.c3_1(x_in[-1])
        x3_out = self.conv2d_1(x)

        x = self.upsample_1(x)
        x = self.concat_1([x, x_in[-2]])
        x = self.conv_2(x)
        x = self.c3_2(x)
        x2_out = self.conv2d_2(x)

        x = self.upsample_2(x)
        x = self.concat_2([x, x_in[-3]])
        x = self.conv_3(x)
        x = self.c3_3(x)
        x1_out = self.conv2d_3(x)

        return [x1_out, x2_out, x3_out]
    

class Detect(nn.Module):
    def __init__(self, nc=80, anchors=()):  # detection layer
        super(Detect, self).__init__()
        self.stride = None  # strides computed during build
        self.nc = nc  # number of classes
        self.no = nc + 5  # number of outputs per anchor
        self.nl = len(anchors)  # number of detection layers
        self.na = len(anchors[0]) // 2  # number of anchors
        self.grid = [torch.zeros(1)] * self.nl  # init grid
        a = torch.tensor(anchors).float().view(self.nl, -1, 2)
        self.register_buffer('anchors', a)  # shape(nl,na,2)
        self.register_buffer('anchor_grid', a.clone().view(self.nl, 1, -1, 1, 1, 2))  # shape(nl,1,na,1,1,2)
        self.export = False  # onnx export

        self. f = [-1,16,11]
        self. i = 22

    def forward(self, x):
        # x = x.copy()  # for profiling
        z = []  # inference output
        self.training |= self.export
        for i in range(self.nl):
            bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
            x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

            if not self.training:  # inference
                if self.grid[i].shape[2:4] != x[i].shape[2:4]:
                    self.grid[i] = self._make_grid(nx, ny).to(x[i].device)

                y = x[i].sigmoid()
                y[..., 0:2] = (y[..., 0:2] * 2. - 0.5 + self.grid[i].to(x[i].device)) * self.stride[i]  # xy
                y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
                z.append(y.view(bs, -1, self.no))

        return x if self.training else (torch.cat(z, 1), x)

    @staticmethod
    def _make_grid(nx=20, ny=20):
        yv, xv = torch.meshgrid([torch.arange(ny), torch.arange(nx)])
        return torch.stack((xv, yv), 2).view((1, 1, ny, nx, 2)).float()
